package com.ihr360.job.core.item;

import com.ihr360.job.core.repeat.RepeatContext;
import com.ihr360.job.core.repeat.exception.ExceptionHandler;
import com.ihr360.job.core.repeat.support.RepeatSynchronizationManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.classify.BinaryExceptionClassifier;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryPolicy;
import org.springframework.retry.listener.RetryListenerSupport;

import java.util.Collection;

/**
 * An {@link ExceptionHandler} that is aware of the retry context so that it can
 * distinguish between a fatal exception and one that can be retried. Delegates
 * the actual exception handling to another {@link ExceptionHandler}.
 *
 * @author Dave Syer
 *
 */
public class SimpleRetryExceptionHandler extends RetryListenerSupport implements ExceptionHandler {

    /**
     * Attribute key, whose existence signals an exhausted retry.
     */
    private static final String EXHAUSTED = SimpleRetryExceptionHandler.class.getName() + ".RETRY_EXHAUSTED";

    private static final Logger logger = LoggerFactory.getLogger(SimpleRetryExceptionHandler.class.getName());

    final private RetryPolicy retryPolicy;

    final private ExceptionHandler exceptionHandler;

    final private BinaryExceptionClassifier fatalExceptionClassifier;

    /**
     * Create an exception handler from its mandatory properties.
     *
     * @param retryPolicy the retry policy that will be under effect when an
     * exception is encountered
     * @param exceptionHandler the delegate to use if an exception actually
     * needs to be handled
     * @param fatalExceptionClasses exceptions
     */
    public SimpleRetryExceptionHandler(RetryPolicy retryPolicy, ExceptionHandler exceptionHandler, Collection<Class<? extends Throwable>> fatalExceptionClasses) {
        this.retryPolicy = retryPolicy;
        this.exceptionHandler = exceptionHandler;
        this.fatalExceptionClassifier = new BinaryExceptionClassifier(fatalExceptionClasses);
    }

    /**
     * Check if the exception is going to be retried, and veto the handling if
     * it is. If retry is exhausted or the exception is on the fatal list, then
     * handle using the delegate.
     *
     * @see ExceptionHandler#handleException(org.springframework.batch.repeat.RepeatContext,
     * java.lang.Throwable)
     */
    @Override
    public void handleException(RepeatContext context, Throwable throwable) throws Throwable {
        // Only bother to check the delegate exception handler if we know that
        // retry is exhausted
        if (fatalExceptionClassifier.classify(throwable) || context.hasAttribute(EXHAUSTED)) {
            logger.debug("Handled fatal exception");
            exceptionHandler.handleException(context, throwable);
        }
        else {
            logger.debug("Handled non-fatal exception", throwable);
        }
    }

    /**
     * If retry is exhausted set up some state in the context that can be used
     * to signal that the exception should be handled.
     *
     * @see org.springframework.retry.RetryListener#close(org.springframework.retry.RetryContext,
     * org.springframework.retry.RetryCallback, java.lang.Throwable)
     */
    @Override
    public <T, E extends Throwable> void close(RetryContext context, RetryCallback<T, E> callback, Throwable throwable) {
        if (!retryPolicy.canRetry(context)) {
            if (logger.isDebugEnabled()) {
                logger.debug("Marking retry as exhausted: "+context);
            }
            getRepeatContext().setAttribute(EXHAUSTED, "true");
        }
    }

    /**
     * Get the parent context (the retry is in an inner "chunk" loop and we want
     * the exception to be handled at the outer "step" level).
     * @return the {@link RepeatContext} that should hold the exhausted flag.
     */
    private RepeatContext getRepeatContext() {
        RepeatContext context = RepeatSynchronizationManager.getContext();
        if (context.getParent() != null) {
            return context.getParent();
        }
        return context;
    }

}
